#include <iostream>
#include <fstream>
#include <cstring>
#include <vector>
#include <queue>
#include "tplate.h"
#include "ti_3D.h"
#include "taxonomy_tree.h"

using namespace std;

taxonomy_tree *init_taxonomy_tree(int tpsize, int NR_cate)
{
    taxonomy_tree *t_tree = new taxonomy_tree[tpsize];
    int max_node = (1 << NR_cate) - 1;
    for (int t = 0; t < tpsize; t++)
    {
        (t_tree + t)->NR_cate = NR_cate;
        (t_tree + t)->node = new one_node[max_node];
        for (int i = 0; i < max_node; i++)
        {
            ((t_tree + t)->node + i)->stat = new float[NR_cate];
            memset(t_tree[t].node[i].stat, 0, sizeof(float) * NR_cate);
        }
    }
    return t_tree;
}

taxonomy_tree *build_taxonomy_tree_3D(train_grid_3D * ti, tplate * tp,
                                      int g)
{
    g=0;
    int cur_cate[ti->NR_cate];
    int neigh_cate[ti->NR_cate];
    taxonomy_tree *t = init_taxonomy_tree(tp->tpsize, ti->NR_cate);
    for (int z = 0; z < ti->height; z++)
    {
        for (int y = 0; y < ti->width; y++)
        {
            for (int x = 0; x < ti->length; x++)
            {
                // FOR each Pixel DO
                memset(cur_cate, 0, ti->NR_cate * sizeof(int));
                cur_cate[ti->data[z * ti->width * ti->length +
                    y * ti->length + x]] = 1;
                for (int i = 0; i < tp->tpsize; i++)
                {
                    // if the neighbor is outside of the image
                    int cur_x = x + tp->pos[i * 3];
                    int cur_y = y + tp->pos[i * 3 + 1];
                    int cur_z = z + tp->pos[i * 3 + 2];
                    if ((cur_x < 0) || (cur_x >= ti->length))
                        continue;
                    if ((cur_y < 0) || (cur_y >= ti->width))
                        continue;
                    if ((cur_z < 0) || (cur_z >= ti->height))
                        continue;
                    // get the categories of current and neighbor cells
                    memset(neigh_cate, 0, ti->NR_cate * sizeof(int));
                    neigh_cate[ti->data
                        [cur_z * ti->width * ti->length +
                        cur_y * ti->length + cur_x]] = 1;
                    update_taxonomy_tree(t, cur_cate, neigh_cate, i);
                }
            }
        }
    }
    return t;
}

void update_taxonomy_tree(taxonomy_tree * t_tree, int *cur_value,
        int *neigh_value, int tp_position)
{
    int cur_node_id = 1;
    for (int c = 0; c < t_tree->NR_cate; c++)
    {
        int have_value = (neigh_value[c] > 0);
        if (have_value > 0)
        {
            for (int i = 0; i < t_tree->NR_cate; i++)
            {
                t_tree[tp_position].node[cur_node_id - 1].stat[i] +=
                    cur_value[i];
            }
        }
        cur_node_id = 2 * cur_node_id + have_value;
    }
}

void search_taxomony_tree(float *result, taxonomy_tree * t_tree,
        int tp_position, int *cond, float *marginal)
{
    memcpy(result, marginal, sizeof(float) * t_tree->NR_cate);
    queue < int >candidate;
    int end_pos = -1;
    for (int c = t_tree->NR_cate - 1; c >= 0; c--)
    {
        if (cond[c] == 1)
        {
            end_pos = c;
            break;
        }
    }
    if (end_pos == -1)
        return;
    candidate.push(1);
    std::vector < int >temp;
    for (int c = 0; c < end_pos; c++)
    {
        temp.clear();
        while (!candidate.empty())
        {
            int cur_value = candidate.front();
            temp.push_back(cur_value * 2);
            if (cond[c] == 0)
                temp.push_back(cur_value * 2 + 1);
            candidate.pop();
        }
        for (size_t i = 0; i < temp.size(); i++)
        {
            candidate.push(temp[i]);
        }
    }
    taxonomy_tree *cur_tree = t_tree + tp_position;
    while (!candidate.empty())
    {
        int cur_value = candidate.front() - 1;
        candidate.pop();
        for (int c = 0; c < t_tree->NR_cate; c++)
        {
            result[c] += cur_tree->node[cur_value].stat[c] + 1;
        }
    }
    int ressum = 0;
    for (int c = 0; c < t_tree->NR_cate; c++)
    {
        ressum += result[c];
    }
    if (ressum < 0)
    {
        memcpy(result, marginal, sizeof(float) * t_tree->NR_cate);
    }
    return;
}

float *extract_all_pdf(data_event * dv, taxonomy_tree * t_tree,
        float *marginal, int NR_cate)
{
    float *result = new float[dv->dvsize * NR_cate];
    for (int i = 0; i < dv->dvsize; ++i)
    {
        search_taxomony_tree(result + i * NR_cate, t_tree,
                dv->tpos[i], dv->value[i], marginal);
    }
    return result;
}

void save_taxonomy_tree(ALL_taxonomy_tree * alltree, int tpsize,
        char *filename)
{
    ofstream ofs;
    ofs.open(filename);
    ofs << alltree->max_scale << endl;
    for (int g = 0; g < alltree->max_scale; g++)
    {
        taxonomy_tree *cur_tree = alltree->t_tree[g];
        ofs << cur_tree->NR_cate << endl;
        int max_node = (1 << cur_tree->NR_cate) - 1;
        for (int t = 0; t < tpsize; t++)
        {
            for (int i = 0; i < max_node; ++i)
            {
                for (int j = 0; j < cur_tree->NR_cate; ++j)
                {
                    ofs << cur_tree[t].node[i].stat[j] << " ";
                }
                ofs << endl;
            }
        }
        ofs << endl;
    }
    ofs.close();
}

void load_taxonomy_tree(ALL_taxonomy_tree * alltree, int tpsize,
        char *filename)
{
    ifstream ifs;
    ifs.open(filename);
    ifs >> alltree->max_scale;
    alltree->t_tree = new taxonomy_tree *[alltree->max_scale];
    for (int g = 0; g < alltree->max_scale; g++)
    {
        alltree->t_tree[g] = new taxonomy_tree[tpsize];
        taxonomy_tree *cur_tree = alltree->t_tree[g];
        ifs >> cur_tree[0].NR_cate;
        int max_node = (1 << cur_tree[0].NR_cate) - 1;
        for (int t = 0; t < tpsize; t++)
        {
            cur_tree[t].NR_cate = cur_tree[0].NR_cate;
            cur_tree[t].node = new one_node[max_node];
            for (int i = 0; i < max_node; ++i)
            {
                cur_tree[t].node[i].stat = new float[cur_tree[t].NR_cate];
                for (int j = 0; j < cur_tree[t].NR_cate; ++j)
                {
                    ifs >> cur_tree[t].node[i].stat[j];
                }
            }
        }
    }
    ifs.close();
}
